'''
Author: duliang thinktanker@163.com
Date: 2025-08-30 22:11:35
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-09-03 00:10:38
FilePath: 
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
# import mss
# import numpy as np
import cv2
from src.window_capture_manager import WindowCapture
import sys
from PyQt5.QtWidgets import QApplication
from findRect import *
# 添加YOLO导入
from myyolo import MyYOLO
import numpy as np

CONF = 0.2
IOU = 0.2
# 创建QApplication实例（PyQt应用需要）
app = QApplication(sys.argv)

# 创建WindowCapture实例
capturer = WindowCapture()

# 获取进程列表
processes = capturer.list_processes()
if not processes:
    print("未找到可用进程")
    exit()

# 显示进程列表供选择
print("可用进程列表:")
for i, proc in enumerate(processes):
    print(f"{i+1}. {proc['name']} ({proc['title']}) - PID: {proc['pid']}")

# 选择进程
try:
    choice = int(input("请选择要捕获的进程编号: ")) - 1
    if 0 <= choice < len(processes):
        selected_proc = processes[choice]
        print(f"选择的进程: {selected_proc['name']} (PID: {selected_proc['pid']})")
    else:
        print("无效的选择")
        exit()
except ValueError:
    print("请输入有效的数字")
    exit()

# 查找窗口
if not capturer.find_window_by_pid_and_title(selected_proc['pid'],
                                             selected_proc['title']):
    print("无法找到选定的进程窗口")
    exit()

print("开始捕获进程窗口，按'q'键退出")
frame = capturer.capture()
cv2.imwrite('screenshot.png', frame)
rect_result = find_rect(frame)

# 初始化YOLO模型
yolo = MyYOLO()
model_path = r"./models/best_0830_640.pt"  # 根据实际情况修改模型路径
if yolo.load_model(model_path):
    print("YOLO模型加载成功")

# 开始捕获并显示
while True:
    # 捕获窗口画面
    frame = capturer.capture()
    if frame is not None and frame.size > 0:
        # 将BGR格式转换为RGB格式以正确显示颜色
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGRA2RGB)
        # 如果找到了矩形区域，则只截取该区域进行显示
        if rect_result:
            x, y, w, h = rect_result[0]
            cropped_frame = frame_rgb[y:y + h, x:x + w]

            # 使用YOLO进行实时检测
            # 转换为BGR以适应YOLO模型
            cropped_bgr = cv2.cvtColor(cropped_frame, cv2.COLOR_RGB2BGR)
            # 创建临时图像进行检测
            temp_image = np.array(cropped_bgr)
            # 执行推理
            if yolo.model is not None:
                results = yolo.model(temp_image, conf=CONF, iou=IOU)
                # 在图片上绘制检测结果
                annotated_image = yolo.annotate_image(cropped_frame.copy(),
                                                      results)
                # 调整图像大小以适应窗口
                window_height, window_width = 600, 800
                resized_image = cv2.resize(annotated_image,
                                           (window_width, window_height))
                cv2.imshow('Process Window Capture', resized_image)
            # 设置窗口置顶
            cv2.setWindowProperty('Process Window Capture',
                                  cv2.WND_PROP_TOPMOST, 1)
            # 设置窗口大小为800x600
            cv2.resizeWindow('Process Window Capture', 800, 600)
            cv2.imwrite('screenshot1.png', cropped_bgr)
        # 按'q'键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        # 添加短暂延迟避免CPU占用过高
        import time
        time.sleep(0.03)

# 关闭所有OpenCV窗口
cv2.destroyAllWindows()
